package com.yiwenup.struct._02_tree.bst.bbst;

import com.yiwenup.struct._02_tree.bst.bbst.base.BalancedBinarySearchTree;

import java.util.Comparator;

/**
 * AVL树：基于二叉搜索树的自平衡
 **/
public class AVLTree<E> extends BalancedBinarySearchTree<E> {
    public AVLTree() {
        this(null);
    }

    public AVLTree(Comparator<E> comparator) {
        super(comparator);
    }

    @Override
    protected void afterAdd(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                // 更新高度
                updateHeight(node);
            } else {
                // 恢复平衡，更新高度
//                reBalance(node);
                reBalance2(node);
                // 整棵树恢复平衡
                break;
            }
        }
    }

    @Override
    protected void afterRemove(Node<E> node) {
        while ((node = node.parent) != null) {
            if (isBalanced(node)) {
                updateHeight(node);
            } else {
                reBalance2(node);
            }
        }
    }

    private boolean isBalanced(Node<E> node) {
        return Math.abs(((AVLNode<E>) node).balanceFactor()) <= 1;
    }

    private void updateHeight(Node<E> node) {
        AVLNode<E> avlNode = (AVLNode<E>) node;
        avlNode.updateHeight();
    }

    private void reBalance(Node<E> g) {
        Node<E> p = ((AVLNode<E>) g).tallerChild();
        Node<E> n = ((AVLNode<E>) p).tallerChild();
        if (p.isLeftChild()) {
            if (n.isLeftChild()) {
                // LL
                rotateRight(g);
            } else {
                // LR
                rotateLeft(p);
                rotateRight(g);
            }
        } else {
            if (n.isRightChild()) {
                // RR
                rotateLeft(g);
            } else {
                // RL
                rotateRight(p);
                rotateLeft(g);
            }
        }
    }

    private void reBalance2(Node<E> g) {
        Node<E> p = ((AVLNode<E>) g).tallerChild();
        Node<E> n = ((AVLNode<E>) p).tallerChild();
        if (p.isLeftChild()) {
            if (n.isLeftChild()) {
                // LL
                rotate(g, n.left, n, n.right, p, p.right, g, g.right);
            } else {
                // LR
                rotate(g, p.left, p, n.left, n, n.right, g, g.right);
            }
        } else {
            if (n.isRightChild()) {
                // RR
                rotate(g, g.left, g, p.left, p, n.left, n, n.right);
            } else {
                // RL
                rotate(g, g.left, g, n.left, n, n.right, p, p.right);
            }
        }
    }

    /**
     * Node r 失衡子树的根节点
     * Node a, b, c 旋转后将成为新的子树的左子树
     * Node d 旋转后将成为新子树的根节点
     * Node e, f, g 旋转后将成为新的子树的右子树
     */
    @Override
    protected void rotate(Node<E> r,
                          Node<E> a, Node<E> b, Node<E> c,
                          Node<E> d,
                          Node<E> e, Node<E> f, Node<E> g) {
        super.rotate(r, a, b, c, d, e, f, g);
        updateHeight(b);
        updateHeight(f);
        updateHeight(d);
    }

    @Override
    protected void afterRotate(Node<E> g, Node<E> p, Node<E> t) {
        super.afterRotate(g, p, t);
        // 高度更新
        updateHeight(g);
        updateHeight(p);
    }

    private static class AVLNode<E> extends Node<E> {
        private int height = 1;

        public AVLNode(E element, Node<E> parent) {
            super(element, parent);
        }

        /**
         * 平衡因子：左子树高度 - 右子树高度
         */
        public int balanceFactor() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            return leftHeight - rightHeight;
        }

        public void updateHeight() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            height = 1 + Math.max(leftHeight, rightHeight);
        }

        public Node<E> tallerChild() {
            int leftHeight = left == null ? 0 : ((AVLNode<E>) left).height;
            int rightHeight = right == null ? 0 : ((AVLNode<E>) right).height;
            if (leftHeight > rightHeight) return left;
            if (leftHeight < rightHeight) return right;
            return isLeftChild() ? left : right;
        }
    }

    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new AVLNode<E>(element, parent);
    }
}
